import os
from struct import Struct
from cStringIO import StringIO


class ListFile(object):
    HEADER = Struct("=L")
    __slots__ = ["file"]
    
    def __init__(self, file):
        self.file = file
    
    @classmethod
    def open(cls, filename):
        return cls(open(filename, "rb"))
    
    def __iter__(self):
        while True:
            data = self._read_item()
            if not data:
                break
            yield data
    
    def _read_item(self):
        header = self.file.read(self.HEADER.size)
        length, = self.HEADER.unpack(header)
        if length == 0:
            return ""
        data = self.file.read(length)
        if len(data) != length:
            raise ValueError("expected %d bytes, found %d" % (length, len(data)))
        return data
    
    def get(self, offset):
        self.file.seek(offset)
        return _read_item()


class CorruptedNodeHeader(Exception):
    pass

class TreeNode(object):
    HEADER = Struct("=QHH")
    __slots__ = ["_file", "_first_child", "_end_of_children", "_end_of_node",
        "data", "has_children"]
    
    def __init__(self, file, start_offset):
        file.seek(start_offset)
        header = file.read(self.HEADER.size)
        if len(header) < self.HEADER.size:
            raise CorruptedNodeHeader(header)
        end_offset, enter_size, leave_size = self.HEADER.unpack(header)
        if enter_size > 0:
            enter_blob = file.read(enter_size)
        else:
            enter_blob = ""
        if leave_size > 0:
            file.seek(end_offset)
            leave_blob = file.read(leave_size)
        else:
            leave_blob = ""
        
        self._file = file
        if end_offset == 0xffffffffffffffff:
            end_offset = os.fstat(file.fileno()).st_size
        self._end_of_children = end_offset
        self._first_child = start_offset + self.HEADER.size + enter_size
        self._end_of_node = end_offset + leave_size 
        self.data = enter_blob + leave_blob
        self.has_children = bool(end_offset - self._first_child)
    
    def __repr__(self):
        return "Node(%r)" % (self.data.encode("hex"),)
    
    def __iter__(self):
        next_child =  self._first_child
        while next_child < self._end_of_children:
            try:
                node = TreeNode(self._file, next_child)
            except CorruptedNodeHeader:
                break
            if node._end_of_children == 0:
                break
            next_child = node._end_of_node
            yield node
    
    def show(self, level = 0):
        print "%s%s" % ("    " * level, self)
        for child in self:
            child.show(level + 1)


class Forest(object):
    def __init__(self, file):
        self.file = file
    
    @classmethod
    def open(cls, filename):
        return cls(open(filename, "rb"))
    
    def __iter__(self):
        offset = 0
        while True:
            try:
                node = TreeNode(self.file, offset)
            except CorruptedNodeHeader:
                break
            if node._end_of_children == 0:
                break
            offset = node._end_of_node
            yield node
    
    def show(self):
        print "Forest"
        for node in self:
            node.show(1)


class StructReader(object):
    __slots__ = ["stream"]
    UI8 = Struct("B")
    UI16 = Struct("H")
    UI32 = Struct("L")
    UI64 = Struct("Q")
    
    def __init__(self, data):
        self.stream = StringIO(data)
    def read(self, count):
        data = self.stream.read(count)
        if len(data) < count:
            raise ValueError("expected %d bytes, found %d" % (count, len(data)))
        return data
    def read_ui8(self):
        return self.UI8.unpack(self.read(1))[0]
    def read_ui16(self):
        return self.UI16.unpack(self.read(2))[0]
    def read_ui32(self):
        return self.UI32.unpack(self.read(4))[0]
    def read_ui64(self):
        return self.UI64.unpack(self.read(8))[0]
    def read_pstr(self):
        return self.read(self.read_ui16())


if __name__ == "__main__":
    f = Forest.open("../../build/lib.linux-i686-2.5/foo-tree")
    f.show()







